package me.zhangjin.lucene.demo.highlighter;

import java.io.IOException;
import java.nio.file.Path;
import java.nio.file.Paths;

import me.zhangjin.lucene.demo.ik.IKAnalyzer6x;
import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.document.Document;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.queryparser.classic.ParseException;
import org.apache.lucene.queryparser.classic.QueryParser;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.search.highlight.Fragmenter;
import org.apache.lucene.search.highlight.Highlighter;
import org.apache.lucene.search.highlight.InvalidTokenOffsetsException;
import org.apache.lucene.search.highlight.QueryScorer;
import org.apache.lucene.search.highlight.SimpleHTMLFormatter;
import org.apache.lucene.search.highlight.SimpleSpanFragmenter;
import org.apache.lucene.search.highlight.TokenSources;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.FSDirectory;


public class HighlighterTest {
    public static void main(String[] args) throws IOException, InvalidTokenOffsetsException, ParseException {

        Path indexPath = Paths.get("indexdir/index");
        Directory dir = FSDirectory.open(indexPath);
        IndexReader reader = DirectoryReader.open(dir);
        IndexSearcher searcher = new IndexSearcher(reader);
        Analyzer analyzer = new IKAnalyzer6x();
        QueryParser parser = new QueryParser("title", analyzer);

        // 使用自定义的分词器构建一个 Query
        Query query = parser.parse("北大");

        // Query:title:北大
        System.out.println("Query:" + query);

        // 返回前10条
        TopDocs tds = searcher.search(query, 10);

        for (ScoreDoc sd : tds.scoreDocs) {

            Document doc = searcher.doc(sd.doc);
            System.out.println("id:" + doc.get("id"));
            System.out.println("title:" + doc.get("title"));

            // 通过 Query
            // 构造一个 QueryScorer，指定了 title 字段用于高亮
            QueryScorer score = new QueryScorer(query, "title");

            // 指定高亮的格式化器（模板）
            SimpleHTMLFormatter fors = new SimpleHTMLFormatter("<span style=\"color:red;\">", "</span>");// 定制高亮标签

            // 通过 QueryScorer 和 SimpleHTMLFormatter
            // 构建一个高亮器
            Highlighter highlighter = new Highlighter(fors, score);

            // 通过 QueryScorer
            // 构建一个片段
            Fragmenter fragment = new SimpleSpanFragmenter(score);

            highlighter.setTextFragmenter(fragment);

            // 获取 TokenStream
            TokenStream tokenStream = TokenSources.getAnyTokenStream(searcher.getIndexReader(), sd.doc, "title", analyzer);

            // 要从 TokenStream 中获取高亮字符串
            // 获取高亮的片段
            String str = highlighter.getBestFragment(tokenStream, doc.get("title"));
            System.out.println("高亮的片段:" + str);
        }
        dir.close();
        reader.close();
    }
}
